import os
import pickle
import wave
import numpy as np
import opuslib_next
from pydub import AudioSegment


class Opus_Encode:

    _instance = None
    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super().__new__(cls)
        return cls._instance
    def __init__(self):
        self.channel = 1
        self.sample_rate = 16000
        self.sample_width = 2
        self.opus_sample_width = 2
        self.opus_channel = 1
        self.opus_sample_rate = 16000
        self.opus_frame_time = 60
        self.frame_size = int(self.opus_sample_rate / 1000 * self.opus_frame_time)

    def audio_to_opus(self, audio_path):
        file_type = os.path.splitext(audio_path)[1]
        print(file_type)

        if file_type:
            file_type = file_type.lstrip(".")

        audio = AudioSegment.from_file(audio_path,format=file_type)
        audio = audio.set_channels(self.channel).set_frame_rate(self.sample_rate).set_sample_width(self.sample_width)
        duration = len(audio) / 1000


        raw_data = audio.raw_data

        print(len(raw_data))
        encoder = opuslib_next.Encoder(self.opus_sample_rate, self.opus_channel, opuslib_next.APPLICATION_AUDIO)
        opus_datas = []
        frame_num = int(self.opus_sample_rate / 1000 * self.opus_frame_time)
        frame_size = frame_num * self.opus_channel * self.opus_sample_width
        for i in range(0, len(raw_data), frame_size):
            chunk = raw_data[i:i + frame_size]
            chunk_len = len(chunk)
            if chunk_len < frame_size:
                chunk += b'\x00' * (frame_size - chunk_len)
            np_frame = np.frombuffer(chunk, dtype=np.int16)
            np_bytes = np_frame.tobytes()
            encoded = encoder.encode(np_bytes,frame_num)
            opus_datas.append(encoded)
        return opus_datas , duration

    def save_opus(self, opus_datas, output_path):
        with open(output_path, "wb") as f:
            pickle.dump(opus_datas, f)
        print("save opus success")
        return output_path

    def load_opus(self, input_path):
        with open(input_path, "rb") as f:
            opus_datas = pickle.load(f)
        return opus_datas

    def opus_to_audio(self, output_path, opus_datas):

        decoder = opuslib_next.Decoder(self.opus_sample_rate,self.channel)
        pcm_data = []
        for frame in opus_datas:
            pcm_frame = decoder.decode(frame,self.frame_size)
            pcm_data.append(pcm_frame)
        with wave.open(output_path,"wb") as f:
            f.setnchannels(self.channel)
            f.setsampwidth(self.sample_width)
            f.setframerate(self.sample_rate)
            f.writeframes(b"".join(pcm_data))
        print("save wav success")
        return output_path

if __name__ == '__main__':
    opus = Opus_Encode()
    opus_datas, duration = opus.audio_to_opus("test.mp3")
    opus.save_opus(opus_datas, duration, "test.opus")
    opus_data = opus.load_opus("test.opus")



